package trace

import (
	"bytes"
	"fmt"
	"math/rand"
	"time"

	"github.com/go-logr/logr"
)

// Field is a key value pair that provides additional details about the trace.
type Field struct {
	Key   string
	Value interface{}
}

func (f Field) format() string {
	return fmt.Sprintf("%s:%v", f.Key, f.Value)
}

func writeFields(b *bytes.Buffer, l []Field) {
	for i, f := range l {
		b.WriteString(f.format())
		if i < len(l)-1 {
			b.WriteString(",")
		}
	}
}

func writeTraceItemSummary(b *bytes.Buffer, msg string, totalTime time.Duration, startTime time.Time, fields []Field) {
	b.WriteString(fmt.Sprintf("%q ", msg))
	if len(fields) > 0 {
		writeFields(b, fields)
		b.WriteString(" ")
	}

	b.WriteString(fmt.Sprintf("%vms (%v)", durationToMilliseconds(totalTime), startTime.Format("15:04:00.000")))
}

func durationToMilliseconds(timeDuration time.Duration) int64 {
	return timeDuration.Nanoseconds() / 1e6
}

type traceItem interface {
	// time returns when the trace was recorded as completed.
	time() time.Time
	// writeItem outputs the traceItem to the buffer. If stepThreshold is non-nil, only output the
	// traceItem if its the duration exceeds the stepThreshold.
	// Each line of output is prefixed by formatter to visually indent nested items.
	writeItem(b *bytes.Buffer, formatter string, startTime time.Time, stepThreshold *time.Duration)
}

type traceStep struct {
	stepTime time.Time
	msg      string
	fields   []Field
}

func (s traceStep) time() time.Time {
	return s.stepTime
}

func (s traceStep) writeItem(b *bytes.Buffer, formatter string, startTime time.Time, stepThreshold *time.Duration) {
	stepDuration := s.stepTime.Sub(startTime)
	if stepThreshold == nil || *stepThreshold == 0 || stepDuration >= *stepThreshold {
		b.WriteString(fmt.Sprintf("%s---", formatter))
		writeTraceItemSummary(b, s.msg, stepDuration, s.stepTime, s.fields)
	}
}

// Trace keeps track of a set of "steps" and allows us to log a specific
// step if it took longer than its share of the total allowed time
type Trace struct {
	name        string
	fields      []Field
	threshold   *time.Duration
	startTime   time.Time
	endTime     *time.Time
	traceItems  []traceItem
	parentTrace *Trace
	logger      logr.Logger
}

func (t *Trace) time() time.Time {
	if t.endTime != nil {
		return *t.endTime
	}
	return t.startTime // if the trace is incomplete, don't assume an end time
}

func (t *Trace) writeItem(b *bytes.Buffer, formatter string, startTime time.Time, stepThreshold *time.Duration) {
	if t.durationIsWithinThreshold() || t.logger.V(4).Enabled() {
		b.WriteString(fmt.Sprintf("%v[", formatter))
		writeTraceItemSummary(b, t.name, t.TotalTime(), t.startTime, t.fields)
		if st := t.calculateStepThreshold(); st != nil {
			stepThreshold = st
		}
		t.writeTraceSteps(b, formatter+" ", stepThreshold)
		b.WriteString("]")
		return
	}
	// If the trace should not be written, still check for nested traces that should be written
	for _, s := range t.traceItems {
		if nestedTrace, ok := s.(*Trace); ok {
			nestedTrace.writeItem(b, formatter, startTime, stepThreshold)
		}
	}
}

// New creates a Trace with the specified name. The name identifies the operation to be traced. The
// Fields add key value pairs to provide additional details about the trace, such as operation inputs.
func New(name string, logger logr.Logger, fields ...Field) *Trace {
	return &Trace{name: name, startTime: time.Now(), fields: fields, logger: logger}
}

// Step adds a new step with a specific message. Call this at the end of an execution step to record
// how long it took. The Fields add key value pairs to provide additional details about the trace
// step.
func (t *Trace) Step(msg string, fields ...Field) {
	if t.traceItems == nil {
		// traces almost always have less than 6 steps, do this to avoid more than a single allocation
		t.traceItems = make([]traceItem, 0, 6)
	}
	t.traceItems = append(t.traceItems, traceStep{stepTime: time.Now(), msg: msg, fields: fields})
}

// Nest adds a nested trace with the given message and fields and returns it.
// As a convenience, if the receiver is nil, returns a top level trace. This allows
// one to call FromContext(ctx).Nest without having to check if the trace
// in the context is nil.
func (t *Trace) Nest(msg string, fields ...Field) *Trace {
	newTrace := New(msg, t.logger, fields...)
	if t != nil {
		newTrace.parentTrace = t
		t.traceItems = append(t.traceItems, newTrace)
	}
	return newTrace
}

// Log is used to dump all the steps in the Trace. It also logs the nested trace messages using indentation.
// If the Trace is nested it is not immediately logged. Instead, it is logged when the trace it is nested within
// is logged.
func (t *Trace) Log() {
	endTime := time.Now()
	t.endTime = &endTime
	// an explicit logging request should dump all the steps out at the higher level
	if t.parentTrace == nil { // We don't start logging until Log or LogIfLong is called on the root trace
		t.logTrace()
	}
}

// LogIfLong only logs the trace if the duration of the trace exceeds the threshold.
// Only steps that took longer than their share or the given threshold are logged.
// If klog is at verbosity level 4 or higher and the trace took longer than the threshold,
// all substeps and subtraces are logged. Otherwise, only those which took longer than
// their own threshold.
// If the Trace is nested it is not immediately logged. Instead, it is logged when the trace it
// is nested within is logged.
func (t *Trace) LogIfLong(threshold time.Duration) {
	t.threshold = &threshold
	t.Log()
}

// logTopLevelTraces finds all traces in a hierarchy of nested traces that should be logged but do not have any
// parents that will be logged, due to threshold limits, and logs them as top level traces.
func (t *Trace) logTrace() {
	if t.durationIsWithinThreshold() {
		var buffer bytes.Buffer
		traceNum := rand.Int31()

		totalTime := t.endTime.Sub(t.startTime)
		buffer.WriteString(fmt.Sprintf("Trace[%d]: %q ", traceNum, t.name))
		if len(t.fields) > 0 {
			writeFields(&buffer, t.fields)
			buffer.WriteString(" ")
		}

		// if any step took more than it's share of the total allowed time, it deserves a higher log level
		buffer.WriteString(fmt.Sprintf("(%v) (total time: %vms):", t.startTime.Format("02-Jan-2006 15:04:05.000"), totalTime.Milliseconds()))
		stepThreshold := t.calculateStepThreshold()
		t.writeTraceSteps(&buffer, fmt.Sprintf("\nTrace[%d]: ", traceNum), stepThreshold)
		buffer.WriteString(fmt.Sprintf("\nTrace[%d]: [%v] [%v] END\n", traceNum, t.endTime.Sub(t.startTime), totalTime))

		t.logger.Info(buffer.String())
		return
	}

	// If the trace should not be logged, still check if nested traces should be logged
	for _, s := range t.traceItems {
		if nestedTrace, ok := s.(*Trace); ok {
			nestedTrace.logTrace()
		}
	}
}

func (t *Trace) writeTraceSteps(b *bytes.Buffer, formatter string, stepThreshold *time.Duration) {
	lastStepTime := t.startTime
	for _, stepOrTrace := range t.traceItems {
		stepOrTrace.writeItem(b, formatter, lastStepTime, stepThreshold)
		lastStepTime = stepOrTrace.time()
	}
}

func (t *Trace) durationIsWithinThreshold() bool {
	if t.endTime == nil { // we don't assume incomplete traces meet the threshold
		return false
	}
	return t.threshold == nil || *t.threshold == 0 || t.endTime.Sub(t.startTime) >= *t.threshold
}

// TotalTime can be used to figure out how long it took since the Trace was created
func (t *Trace) TotalTime() time.Duration {
	return time.Since(t.startTime)
}

// calculateStepThreshold returns a threshold for the individual steps of a trace, or nil if there is no threshold and
// all steps should be written.
func (t *Trace) calculateStepThreshold() *time.Duration {
	if t.threshold == nil {
		return nil
	}
	lenTrace := len(t.traceItems) + 1
	traceThreshold := *t.threshold
	for _, s := range t.traceItems {
		nestedTrace, ok := s.(*Trace)
		if ok && nestedTrace.threshold != nil {
			traceThreshold = traceThreshold - *nestedTrace.threshold
			lenTrace--
		}
	}

	// the limit threshold is used when the threshold(
	//remaining after subtracting that of the child trace) is getting very close to zero to prevent unnecessary logging
	limitThreshold := *t.threshold / 4
	if traceThreshold < limitThreshold {
		traceThreshold = limitThreshold
		lenTrace = len(t.traceItems) + 1
	}

	stepThreshold := traceThreshold / time.Duration(lenTrace)
	return &stepThreshold
}
